/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
#ifndef IMPL_REDUCE_REDUCE_COMMON_UTIL_IMPL_H
#define IMPL_REDUCE_REDUCE_COMMON_UTIL_IMPL_H

#include "kernel_operator_intf.h"
#include "kernel_tensor.h"
#include "../common/check.h"

namespace AscendC {
namespace Pattern {
namespace Detail {
constexpr int32_t DIM_TWO = 2;
constexpr int32_t DIM_THREE = 3;
constexpr int32_t DIM_FOUR = 4;
constexpr int32_t DIM_FIVE = 5;
constexpr int32_t DIM_SIX = 6;
constexpr int32_t DIM_SEVEN = 7;
constexpr int32_t DIM_EIGHT = 8;
constexpr int32_t DIM_NINE = 9;
constexpr int32_t PATTERN_R = 0;
constexpr int32_t PATTERN_RA = 1;
constexpr int32_t PATTERN_AR = 2;
constexpr int32_t PATTERN_ARA = 3;
constexpr int32_t PATTERN_ARAR = 4;
constexpr int32_t PATTERN_ARARA = 5;
constexpr int32_t PATTERN_ARARAR = 6;
constexpr int32_t PATTERN_ARARARA = 7;
constexpr int32_t PATTERN_ARARARAR = 8;
constexpr int32_t PATTERN_ARARARARA = 9;
constexpr int32_t PATTERN_RAR = 10;
constexpr int32_t PATTERN_RARA = 11;
constexpr int32_t PATTERN_RARAR = 12;
constexpr int32_t PATTERN_RARARA = 13;
constexpr int32_t PATTERN_RARARAR = 14;
constexpr int32_t PATTERN_RARARARA = 15;

template <int32_t id, bool firstA, bool tailA, int32_t dim>
struct PatternConstInfo {
constexpr static int32_t ID = id;
constexpr static bool FirstA = firstA;
constexpr static bool TailA = tailA;
constexpr static int32_t Dim = dim;
};
} // namespace Detail
} // namespace Pattern

namespace Internal {
// api mode
enum class ApiMode : uint8_t {
    API_MODE_SUM = 0,
    API_MODE_MIN,
    API_MODE_MAX,
    API_MODE_ANY,
    API_MODE_ALL
};

// Find the closest power of two, except 0.
__aicore__ inline uint32_t FindClosestPowerOfTwo(uint32_t n)
{
    ASCENDC_ASSERT(n != 0, { KERNEL_LOG(KERNEL_ERROR, "input n must be non-zero!"); });
    constexpr uint32_t totalShiftBits = 63;
    return totalShiftBits - ScalarCountLeadingZero(n);
}

template <class T>
__aicore__ inline void ComputeMaskBit(uint32_t oneBlkMask, uint32_t oneBlkElems, uint32_t blkNum,
    uint64_t& maskLow, uint64_t& maskHigh)
{
    // Get bit mask on each blkNum, oneBlkElems means number of elements in one block
    // only support half/float
    if constexpr (sizeof(T) == sizeof(half)) {
        uint32_t maskLBlkNum = blkNum > HALF_DEFAULT_REPEAT_STRIDE ? HALF_DEFAULT_REPEAT_STRIDE : blkNum;
        uint32_t maskHBlkNum = blkNum - maskLBlkNum;
        for (int32_t k = 0; k < maskLBlkNum; k++) {
            maskLow += (((1ULL << oneBlkMask) - 1ULL) << (k * oneBlkElems));
        }
        for (int32_t k = 0; k < maskHBlkNum; k++) {
            maskHigh += (((1ULL << oneBlkMask) - 1ULL) << (k * oneBlkElems));
        }
    } else if constexpr (sizeof(T) == sizeof(float)) {
        for (int32_t k = 0; k < blkNum; k++) {
            maskLow += (((1ULL << oneBlkMask) - 1ULL) << (k * oneBlkElems));
        }
    }
}

template <class T, ApiMode apiMode, MaskMode maskMode = MaskMode::NORMAL>
__aicore__ inline void BlockReduceCompute(const LocalTensor<T>& dstTensor, const LocalTensor<T>& srcTensor,
    const int32_t repeat, const uint64_t mask[], const int32_t blkStride, const int32_t repStride)
{
    if constexpr (maskMode == MaskMode::NORMAL) {
        if constexpr (apiMode == ApiMode::API_MODE_SUM) {
            BlockReduceSum(dstTensor, srcTensor, repeat, mask, 1, 1, DEFAULT_REPEAT_STRIDE);
        } else if constexpr (apiMode == ApiMode::API_MODE_MIN || apiMode == ApiMode::API_MODE_ALL) {
            BlockReduceMin(dstTensor, srcTensor, repeat, mask, 1, 1, DEFAULT_REPEAT_STRIDE);
        } else if constexpr (apiMode == ApiMode::API_MODE_MAX || apiMode == ApiMode::API_MODE_ANY) {
            BlockReduceMax(dstTensor, srcTensor, repeat, mask, 1, 1, DEFAULT_REPEAT_STRIDE);
        }
    } else {
        if constexpr (apiMode == ApiMode::API_MODE_MIN || apiMode == ApiMode::API_MODE_ALL) {
            BlockReduceMin<T, false>(dstTensor, srcTensor, 1, MASK_PLACEHOLDER, 1, blkStride, repStride);
        } else if constexpr (apiMode == ApiMode::API_MODE_MAX || apiMode == ApiMode::API_MODE_ANY) {
            BlockReduceMax<T, false>(dstTensor, srcTensor, 1, MASK_PLACEHOLDER, 1, blkStride, repStride);
        }
    }
}

template <class T, ApiMode apiMode>
__aicore__ inline void WholeReduceCompute(const LocalTensor<T>& dstTensor, const LocalTensor<T>& srcTensor,
    const int32_t repeat, const int32_t mask, const int32_t repStride)
{
    if constexpr (apiMode == ApiMode::API_MODE_MIN || apiMode == ApiMode::API_MODE_ALL) {
        WholeReduceMin(dstTensor, srcTensor, mask, repeat, 1, 1, repStride, ReduceOrder::ORDER_ONLY_VALUE);
    } else if constexpr (apiMode == ApiMode::API_MODE_MAX || apiMode == ApiMode::API_MODE_ANY) {
        WholeReduceMax(dstTensor, srcTensor, mask, repeat, 1, 1, repStride, ReduceOrder::ORDER_ONLY_VALUE);
    }
}
} // namespace Internal
} // namespace AscendC
#endif // IMPL_REDUCE_REDUCE_COMMON_UTIL_IMPL_H